import gradio as gr
import asyncio
from modules.translation.translator import TranslationModel
from modules.translation.language_detector import LanguageDetector

# 初始化语言检测器
language_detector = LanguageDetector()

async def translate_text(text: str, source_lang: str = "auto", target_lang: str = "zh") -> str:
    """翻译文本"""
    try:
        if not text.strip():
            return "请输入要翻译的文本"
            
        # 自动检测语言
        if source_lang == "auto":
            detected_lang = await language_detector.detect_language(text)
            if not detected_lang:
                return "无法检测输入语言"
            source_lang = detected_lang
            
        # 如果源语言和目标语言相同，直接返回原文
        if source_lang == target_lang:
            return text
            
        # 初始化翻译模型
        model = TranslationModel(
            f"Helsinki-NLP/opus-mt-{source_lang}-{target_lang}",
            config={
                "source_lang": source_lang,
                "target_lang": target_lang,
                "max_length": 512,
                "num_beams": 4
            }
        )
        
        # 加载模型并翻译
        await model.load()
        translation = await model.predict(text)
        return translation
        
    except Exception as e:
        return f"翻译错误: {str(e)}"

def translate_wrapper(text: str, source_lang: str, target_lang: str) -> str:
    """包装异步翻译函数以供 Gradio 使用"""
    return asyncio.run(translate_text(text, source_lang, target_lang))

# 创建 Gradio 界面
with gr.Blocks(title="PolyModel 翻译器") as demo:
    gr.Markdown("# PolyModel 翻译器")
    gr.Markdown("支持中英互译的智能翻译系统")
    
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(
                label="输入文本",
                placeholder="请输入要翻译的文本...",
                lines=5
            )
            
            with gr.Row():
                source_lang = gr.Dropdown(
                    choices=["auto", "en", "zh"],
                    value="auto",
                    label="源语言"
                )
                target_lang = gr.Dropdown(
                    choices=["en", "zh"],
                    value="zh",
                    label="目标语言"
                )
            
            translate_btn = gr.Button("翻译")
        
        with gr.Column():
            output_text = gr.Textbox(
                label="翻译结果",
                lines=5
            )
    
    # 添加示例
    gr.Examples(
        examples=[
            ["Hello, how are you?", "auto", "zh"],
            ["The weather is beautiful today!", "en", "zh"],
            ["你好，最近过得怎么样？", "auto", "en"],
            ["今天天气真不错！", "zh", "en"],
            ["The neural network uses backpropagation to optimize weights.", "en", "zh"],
            ["深度学习模型需要大量训练。", "zh", "en"]
        ],
        inputs=[input_text, source_lang, target_lang],
        outputs=output_text,
        fn=translate_wrapper,
        cache_examples=True
    )
    
    # 绑定翻译按钮事件
    translate_btn.click(
        fn=translate_wrapper,
        inputs=[input_text, source_lang, target_lang],
        outputs=output_text
    )

# 启动应用
if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",  # 允许外部访问
        server_port=7860,
        share=True  # 生成公共链接
    )
